package lil.llahn.security.common.asymmetric;

import lil.llahn.security.common.utils.Base64;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * @ClassName: DHSecurityTools
 * @Description: TODO Diffie-Hellman - 一种确保共享KEY安全穿越不安全网络的方法，也就是常说的密钥一致协议。
 * 简单的说就是允许两名用户在公开媒体上交换信息以生成“一致”的、可以共享的密钥。也就是由甲方产出一对密钥（公钥、私钥），乙方依照甲方公钥产生乙方密钥对（公钥、私钥）。
 * @Date: 2019/11/12 11:07
 * @Auther: lil
 */
public class DHSecurityTools {
    private static final String KEY_DH = "DH";  // 指定加密方式
    private static final int KEY_SIZE = 1024;  // 默认秘钥长度设置
    private static final String KEY_DH_DES = "DES";  // DH加密下需要一种对称加密算法对数据加密，这里我们使用DES，也可以使用其他对称加密算法

    static{
        /*
         * todo 由于JDK8 update 161之后，DH的密钥长度至少为512位，但DES算法密钥不能达到这样的长度，所以导致报错。
         *  Unsupported secret key algorithm：DES
         *  添加以下设置，可以解决问题
         */
        System.setProperty("jdk.crypto.KeyAgreement.legacyKDF", "true");
    }
    /**
     * todo 初始化 甲方 公钥 私钥
     *
     * @return
     */
    public static Map<Integer, String> init() {
        Map<Integer, String> map = new HashMap<>();
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance(KEY_DH);
            generator.initialize(KEY_SIZE);
            KeyPair keyPair = generator.generateKeyPair();
            DHPublicKey publicKey = (DHPublicKey) keyPair.getPublic();
            DHPrivateKey privateKey = (DHPrivateKey) keyPair.getPrivate();
            map.put(0, Base64.encoderToString(publicKey.getEncoded()));  // 公钥
            map.put(1, Base64.encoderToString(privateKey.getEncoded())); // 私钥
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage(), e);
        }
        return map;
    }

    /**
     * todo 初始化 乙方 公钥 私钥
     * @param key  甲方 公钥
     * @return
     */
    public static Map<Integer, String> init(String key) {
        Map<Integer, String> map = new HashMap<>();
        try{
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.decoderToByte(key));
            KeyFactory factory = KeyFactory.getInstance(KEY_DH);
            PublicKey publicKey = factory.generatePublic(keySpec);

            DHParameterSpec spec = ((DHPublicKey) publicKey).getParams();
            KeyPairGenerator generator = KeyPairGenerator.getInstance(KEY_DH);
            generator.initialize(spec);
            KeyPair keyPair = generator.generateKeyPair();
            // 乙方公钥
            DHPublicKey dhPublicKey = (DHPublicKey) keyPair.getPublic();
            // 乙方私钥
            DHPrivateKey dhPrivateKey = (DHPrivateKey) keyPair.getPrivate();
            map.put(0, Base64.encoderToString(dhPublicKey.getEncoded()));  // 公钥
            map.put(1, Base64.encoderToString(dhPrivateKey.getEncoded())); // 私钥
        }catch (Exception e){
            throw new RuntimeException(e.getMessage(), e);
        }
        return map;
    }

    /**
     * todo 构建本地秘钥
     * @param publicKey  本方的公钥
     * @param privateKey 对方的私钥
     * @return
     */
    private static SecretKey getSecretKey(String publicKey, String privateKey){
        SecretKey secretKey = null;
        try{
            KeyFactory factory = KeyFactory.getInstance(KEY_DH);

            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.decoderToByte(publicKey));
            PublicKey localPublicKey = factory.generatePublic(keySpec);

            PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(Base64.decoderToByte(privateKey));
            PrivateKey localPrivateKey = factory.generatePrivate(spec);

            KeyAgreement agreement = KeyAgreement.getInstance(factory.getAlgorithm());
            agreement.init(localPrivateKey);
            agreement.doPhase(localPublicKey, true);

            // 生成本地密钥
            secretKey = agreement.generateSecret(KEY_DH_DES);
        }catch (Exception e){
            throw new RuntimeException(e.getMessage(), e);
        }
        return secretKey;
    }

    /**
     * todo DH 加密
     * @param data  待加密数据
     * @param publicKey  本方公钥
     * @param privateKey 对方私钥
     * @return  加密后的字符串
     */
    public static String encrypt(String data, String publicKey, String privateKey) {
        try{
            // 生成本地密钥
            SecretKey secretKey = getSecretKey(publicKey, privateKey);

            Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, secretKey);
            return Base64.encoderToString(cipher.doFinal(data.getBytes()));
        }catch (Exception e){
            throw new RuntimeException(e.getMessage(), e);
        }
    }

    /**
     * todo DH 解密
     * @param data  待解密字符串
     * @param publicKey  本方公钥
     * @param privateKey 对方私钥
     * @return 解密后的字符串
     */
    public static String decrypt(String data, String publicKey, String privateKey) {
        try{
            // 生成本地密钥
            SecretKey secretKey = getSecretKey(publicKey, privateKey);

            Cipher cipher = Cipher.getInstance(secretKey.getAlgorithm());
            cipher.init(Cipher.DECRYPT_MODE, secretKey);
            return new String(cipher.doFinal(Base64.decoderToByte(data)));
        }catch (Exception e){
            throw new RuntimeException(e.getMessage(), e);
        }
    }
}
